# Copyright 2021 Dakewe Biotech Corporation. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""File description: Realize the function of converting low-resolution image to high-resolution image using the model."""

import os

import cv2
import torch
import numpy as np
import imgproc
from model import EDSR  # 假设 EDSR 是您定义的超分辨率模型

# 定义相关的配置变量
upscale_factor = 2  # 替换为您实际的放大倍数,必须是模型支持的才行，否则加载会报错
# model_path = 'models/edsr_x2-DIV2K-ed69cda3.pth.tar'  # 替换为您的模型路径
device = "cuda" if torch.cuda.is_available() else "cpu"  # 根据您的设备情况选择

# 模型路径
model_path = f"models/edsr_x2-DIV2K-ed69cda3.pth.tar"

def convert_lr_to_hr(lr_image_path, output_path):
    # Initialize the super-resolution model

    model = EDSR(upscale_factor).to(device)
    print("Build EDSR model successfully.")

    print(model_path)
    # Load the super-resolution model weights
    checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)
    model.load_state_dict(checkpoint["state_dict"])
    print(f"Load EDSR model weights `{os.path.abspath(model_path)}` successfully.")

    # Read the low-resolution image
    lr_image = cv2.imread(lr_image_path, cv2.IMREAD_UNCHANGED).astype(np.float32) / 255.0

    # Convert BGR image to RGB image
    lr_image = cv2.cvtColor(lr_image, cv2.COLOR_BGR2RGB)

    # Convert RGB image data to tensor data
    lr_tensor = imgproc.image2tensor(lr_image, range_norm=False, half=True).to(device).unsqueeze_(0)

    model.eval()
    # Turn on half-precision inference.
    model.half()

    # Only reconstruct the Y channel image data.
    with torch.no_grad():
        sr_tensor = model(lr_tensor).clamp_(0, 1.0)

    # Save the high-resolution image
    sr_image = imgproc.tensor2image(sr_tensor, range_norm=False, half=True)
    sr_image = cv2.cvtColor(sr_image, cv2.COLOR_RGB2BGR)
    cv2.imwrite(output_path, sr_image)

if __name__ == "__main__":
    lr_image_path = f'data/Capture001.jpg'  # 替换为您的低分辨率图片路径
    output_path = f'converted_hr_image4.jpg'  # 输出的高分辨率图片保存路径
    convert_lr_to_hr(lr_image_path, output_path)